import wave
import tempfile
import sounddevice as sd
import json
import os

#加載用戶記憶
MEMORY_FILE = "memory.json"
try:
    with open(MEMORY_FILE,  "r", encoding="utf-8") as f:
        memory = json.load(f)
except FileNotFoundError:
    memory = {}

def record_audio():
    """Record audio and save to a temporary WAV file."""
    print("Press Enter to start recording, q and Enter to exit.")
    s = input()
    if s == 'q':
        return ''
    print("Recording... Press Enter again to stop, q and Enter to exit.")

    temp_wav_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
    temp_wav_path = temp_wav_file.name

    fs = 48000  # Sampling frequency
    channels = 1
    dtype = 'int16'  # Match WAV file format
    frames = []

    def callback(indata, frame_count, time_info, status):
        if status:
            print(f"Recording status: {status}", flush=True)
        frames.append(indata.copy())

    with sd.InputStream(samplerate=fs, channels=channels, dtype=dtype, callback=callback):
        s = input()  # Wait for the user to press Enter again to stop recording
        if s == 'q':
            return ''

    # Save the recorded data to a WAV file
    with wave.open(temp_wav_path, 'wb') as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(2)  # 16-bit audio
        wf.setframerate(fs)
        wf.writeframes(b''.join(frames))

    print(f"Recording saved to {temp_wav_path}")
    return temp_wav_path


def main():
    from llm import LlmModel
    llm_model = LlmModel('deepseek-r1-1.5b')

    from asr import AsrModel
    asr_model = AsrModel()

    while True:
        audio_file = record_audio()
        if audio_file == '':
            print('Bye')
            break
        text = asr_model(audio_file)
        print('User:', text)

        #記住指令
        if "記住" in text and "喜歡" in text:
            try:
                key = text.split("喜歡")[1].strip()
                memory[key] = True
                with open(MEMORY_FILE,  "w", encoding="utf-8") as f:
                    json.dump(memory, f, ensure_ascii=False, indent=2)
                print("小景已記住你的喜歡" + key)
                continue
            except Exception as e:
                print("記住失敗：", e)


        cont = llm_model.generate(text)
        print('AI:')
        full_cont = ''
        for c in cont:
            print(c, end='', flush=True)
            full_cont += c
        print()


if __name__ == '__main__':
    main()
